-- 导入所需模块
import Data.Char (isSpace)       -- 用于检查空白字符
import Data.List (find, intercalate)  -- 列表操作函数
import Data.Maybe (fromMaybe)    -- Maybe类型处理
import Text.Read (readMaybe)     -- 安全读取函数
import Control.Monad.State       -- State Monad用于管理环境
import System.IO (hFlush, stdout) -- 控制台输出刷新

-- ===== 类型定义 =====

-- 抽象语法树类型
data AST
  = Atom String   -- 原子类型：符号或数字(如'a', '42')
  | List [AST]    -- 列表类型：包含多个AST节点(如'(a b c)')
  deriving (Show, Eq)  -- 自动派生Show和Eq实例

-- 值类型：求值结果
data Value
  = Symbol String     -- 符号值(如变量名)
  | Number Int        -- 整数值
  | Pair Value Value  -- 点对：car和cdr(用于构建列表)
  | Nil               -- 空列表
  | Bool Bool         -- 布尔值(Lisp中用t和()表示真/假)
  | Lambda [String] AST Env  -- 闭包：参数列表、函数体和定义时的环境
  deriving (Show, Eq)

-- 环境类型：存储变量绑定(变量名到值的映射列表)
type Env = [(String, Value)]

-- ===== 词法分析器 =====

-- 定义Token类型
data Token 
  = LParen       -- 左括号'('
  | RParen       -- 右括号')'
  | AtomToken String  -- 原子token(符号或数字)
  | QuoteToken   -- 引号'\''
  deriving (Show, Eq)

-- 将字符串转换为token列表
tokenize :: String -> [Token]
tokenize = tokenize'  -- 使用辅助函数tokenize'
  where
    -- 递归处理输入字符串
    tokenize' "" = []  -- 空字符串返回空列表
    
    -- 处理左括号
    tokenize' ('(':cs) = LParen : tokenize' (dropWhile isSpace cs)
    
    -- 处理右括号
    tokenize' (')':cs) = RParen : tokenize' (dropWhile isSpace cs)
    
    -- 处理引号
    tokenize' ('\'':cs) = QuoteToken : tokenize' (dropWhile isSpace cs)
    
    -- 处理原子(符号或数字)
    tokenize' cs =
      -- 找到下一个分隔符(空白或括号或引号)
      let (atom, rest) = break (\c -> isSpace c || c `elem` "()'") cs
      -- 创建原子token并继续处理剩余部分
      in AtomToken atom : tokenize' (dropWhile isSpace rest)

-- ===== 语法分析器 =====

-- 主解析函数：将token列表转换为AST
parse :: [Token] -> AST
parse tokens = 
  case parseExpr tokens of
    (ast, []) -> ast  -- 成功解析且无剩余token
    (_, rest) -> error $ "Unexpected tokens after expression: " ++ show rest

-- 解析单个表达式
parseExpr :: [Token] -> (AST, [Token])
parseExpr [] = error "Unexpected end of input"  -- 空输入报错
parseExpr (LParen : tokens) =  -- 处理列表开始
  case parseList tokens [] of
    (list, RParen : rest) -> (List list, rest)  -- 找到匹配的右括号
    _ -> error "Expected ')'"  -- 没有匹配的右括号报错
parseExpr (AtomToken s : tokens) = (Atom s, tokens)  -- 处理原子
parseExpr (QuoteToken : tokens) =  -- 处理引号表达式
  let (quotedAst, rest) = parseExpr tokens  -- 解析被引用的表达式
  in (List [Atom "quote", quotedAst], rest)  -- 转换为(quote ...)形式
parseExpr (RParen : _) = error "Unexpected ')'"  -- 多余的右括号报错

-- 解析列表内容(递归处理)
parseList :: [Token] -> [AST] -> ([AST], [Token])
parseList [] _ = error "Unexpected end of input"  -- 列表未闭合报错
parseList (RParen : tokens) acc = (reverse acc, RParen : tokens)  -- 遇到右括号返回结果
parseList tokens acc =
  let (ast, rest) = parseExpr tokens  -- 解析列表中的表达式
  in parseList rest (ast : acc)  -- 递归处理剩余token

-- ===== 辅助函数 =====

-- 尝试将字符串解析为整数，否则视为符号
tryInt :: String -> Value
tryInt s = maybe (Symbol s) Number (readMaybe s)  -- 使用readMaybe安全解析

-- 检查值是否为真(Lisp真值规则)
isTrue :: Value -> Bool
isTrue (Bool False) = False  -- 显式假值为假
isTrue Nil = False           -- 空列表为假
isTrue _ = True              -- 其他情况为真

-- 将AST直接转换为值(用于quote特殊形式)
astToValue :: AST -> Value
astToValue (Atom s) = tryInt s  -- 原子转换为符号或数字
astToValue (List lst) = buildList lst  -- 列表转换为点对链表
  where
    buildList [] = Nil  -- 空列表转换为Nil
    buildList (x:xs) = Pair (astToValue x) (buildList xs)  -- 递归构建点对

-- 创建初始环境(空环境)
initialEnv :: Env
initialEnv = []

-- ===== 求值器 =====

-- 使用State Monad管理环境
type Eval a = State Env a  -- Eval类型别名，State Monad携带环境

-- 主求值函数：将AST转换为Value
eval :: AST -> Eval Value
eval (Atom s) = do  -- 处理原子
  env <- get  -- 获取当前环境
  case lookup s env of  -- 查找变量
    Just v -> return v  -- 找到则返回值
    Nothing -> return $ tryInt s  -- 未找到则尝试作为数字或符号

eval (List []) = return Nil  -- 空列表求值为Nil

-- 处理quote特殊形式
eval (List [Atom "quote", arg]) = 
  return $ astToValue arg  -- 直接转换参数为值而不求值

-- 处理atom特殊形式(检查参数是否为原子)
eval (List [Atom "atom", arg]) = do
  v <- eval arg  -- 先求值参数
  return $ case v of  -- 检查是否为原子类型
    Symbol _ -> Bool True
    Number _ -> Bool True
    Nil -> Bool True
    _ -> Bool False

-- 处理eq特殊形式(相等性判断)
eval (List [Atom "eq", lhs, rhs]) = do
  lv <- eval lhs  -- 求值左参数
  rv <- eval rhs  -- 求值右参数
  return $ evalEq lv rv  -- 比较值
  where
    -- 递归比较点对
    evalEq (Symbol a) (Symbol b) = Bool (a == b)
    evalEq (Number a) (Number b) = Bool (a == b)
    evalEq Nil Nil = Bool True
    evalEq (Pair a1 b1) (Pair a2 b2) = 
      if evalEq a1 a2 == Bool True 
      then evalEq b1 b2 
      else Bool False
    evalEq _ _ = Bool False

-- 处理car特殊形式(获取点对的第一个元素)
eval (List [Atom "car", arg]) = do
  v <- eval arg  -- 求值参数
  case v of
    Pair car _ -> return car  -- 返回car部分
    Nil -> return Nil         -- 空列表的car是Nil
    _ -> error "car: expected pair"  -- 非点对报错

-- 处理cdr特殊形式(获取点对的第二个元素)
eval (List [Atom "cdr", arg]) = do
  v <- eval arg  -- 求值参数
  case v of
    Pair _ cdr -> return cdr  -- 返回cdr部分
    Nil -> return Nil         -- 空列表的cdr是Nil
    _ -> error "cdr: expected pair"  -- 非点对报错

-- 处理cons特殊形式(构造点对)
eval (List [Atom "cons", lhs, rhs]) = do
  lv <- eval lhs  -- 求值第一个参数
  rv <- eval rhs  -- 求值第二个参数
  return $ Pair lv rv  -- 创建点对

-- 处理cond特殊形式(条件判断)
eval (List (Atom "cond" : clauses)) = 
  evalClauses clauses  -- 处理所有条件子句
  where
    evalClauses [] = error "cond: no true clause"  -- 没有真子句报错
    evalClauses (List [test, result] : rest) = do  -- 处理每个(test result)子句
      testVal <- eval test  -- 求值测试表达式
      if isTrue testVal 
        then eval result    -- 测试为真则求值结果
        else evalClauses rest  -- 否则继续处理剩余子句
    evalClauses _ = error "cond: malformed clause"  -- 格式错误报错

-- 处理lambda特殊形式(创建闭包)
eval (List [Atom "lambda", List params, body]) = do
  paramNames <- mapM getParamName params  -- 获取参数名列表
  closureEnv <- get
  return $ Lambda paramNames body closureEnv  -- 返回闭包(捕获当前环境)
  where
    getParamName (Atom s) = return s  -- 参数必须是原子
    getParamName _ = error "lambda: parameter must be atom"

-- 处理label特殊形式(定义命名函数)
eval (List [Atom "label", Atom name, expr]) = do
  value <- eval expr  -- 求值表达式
  modify ((name, value) :)  -- 将名称绑定到环境中
  return value  -- 返回绑定的值

-- 处理加法运算
eval (List [Atom "+", lhs, rhs]) = do
  l <- eval lhs  -- 求值左操作数
  r <- eval rhs  -- 求值右操作数
  case (l, r) of
    (Number a, Number b) -> return $ Number (a + b)  -- 执行加法
    _ -> error "+ not called with two numbers"  -- 非数字报错

-- 处理减法运算(类似加法)
eval (List [Atom "-", lhs, rhs]) = do
  l <- eval lhs
  r <- eval rhs
  case (l, r) of
    (Number a, Number b) -> return $ Number (a - b)
    _ -> error "- not called with two numbers"

-- 处理乘法运算(类似加法)
eval (List [Atom "*", lhs, rhs]) = do
  l <- eval lhs
  r <- eval rhs
  case (l, r) of
    (Number a, Number b) -> return $ Number (a * b)
    _ -> error "* not called with two numbers"

-- 处理除法运算(类似加法)
eval (List [Atom "/", lhs, rhs]) = do
  l <- eval lhs
  r <- eval rhs
  case (l, r) of
    (Number a, Number b) -> 
      if b == 0 then error "Division by zero"  -- 检查除零错误
      else return $ Number (a `div` b)
    _ -> error "/ not called with two numbers"

-- 处理函数应用(第一个元素是函数，后面是参数)
eval (List (func : args)) = do
  funcVal <- eval func  -- 求值函数
  argVals <- mapM eval args  -- 求值所有参数
  apply funcVal argVals  -- 应用函数

-- 应用函数到参数
apply :: Value -> [Value] -> Eval Value
apply (Lambda params body closureEnv) argVals
  | length params == length argVals = do  -- 检查参数数量
      let newBindings = zip params argVals  -- 创建新绑定
      currentEnv <- get  -- 获取当前环境
      -- 保存当前环境，使用闭包环境作为基础
      put (newBindings ++ closureEnv)  -- 扩展环境
      result <- eval body  -- 在新环境中求值函数体
      -- 恢复环境
      put currentEnv
      return result
  | otherwise = error $ "Wrong number of arguments: expected " ++ 
                show (length params) ++ ", got " ++ show (length argVals)
apply x y = error $ "Cannot apply non-function: " ++ show x ++ show y  -- 非函数应用报错

-- ===== 打印函数 =====

-- 将Value转换为字符串表示
stringOfValue :: Value -> String
stringOfValue (Symbol s) = s  -- 符号直接显示
stringOfValue (Number n) = show n  -- 数字转换为字符串
stringOfValue (Bool True) = "t"  -- Lisp的真值显示为t
stringOfValue (Bool False) = "()"  -- Lisp的假值显示为()
stringOfValue Nil = "()"  -- 空列表显示为()
stringOfValue (Pair car cdr) =  -- 点对显示为列表形式
  "(" ++ stringOfValue car ++ printList cdr ++ ")"
  where
    -- 辅助函数：递归打印列表剩余部分
    printList (Pair hd tl) = " " ++ stringOfValue hd ++ printList tl
    printList Nil = ""  -- 列表结束
    printList v = " . " ++ stringOfValue v  -- 非正规列表显示点对
stringOfValue (Lambda params _ _) = "lambda " ++ unwords params  -- 闭包显示参数

-- ===== 组合函数 =====

-- 完整流程：字符串 -> token列表 -> AST -> 求值
sExprToValue :: String -> State Env Value
sExprToValue s = 
  let tokens = tokenize s  -- 词法分析
      ast = parse tokens   -- 语法分析
  in eval ast              -- 求值

-- ===== 测试函数 =====

testBuiltins :: String
testBuiltins = 
  let tests =  -- 定义测试用例列表
        [ ("(quote a)", "a")
        , ("(quote (a b c))", "(a b c)")
        , ("(atom 'a)", "t")
        , ("(atom '(a b))", "()")
        , ("(eq 'a 'a)", "t")
        , ("(eq 'a 'b)", "()")
        , ("(car '(a b c))", "a")
        , ("(cdr '(a b c))", "(b c)")
        , ("(cons 'a '(b c))", "(a b c)")
        , ("(cond ((eq 1 2) 'a) ((eq 1 1) 'b))", "b")
        , ("()", "()")
        , ("'(1 2 3)", "(1 2 3)")
        , ("'(a b c)", "(a b c)")
        , ("(cons 'a (cons 'b ()))", "(a b)")
        , ("(cond (() 'a) ('t 'b))", "b")
        , ("(+ 3 5)", "8")
        , ("((lambda (x) (cons 'a x)) '(b c))", "(a b c)")
        , ("((lambda (x y) (cons x (cons y ()))) 'a 'b)", "(a b)")
        , ("((lambda (x) (cons (car x) (cdr x))) '(a b c))", "(a b c)")
        , ("(label add3 (lambda (x) (+ x 3)))", "lambda x")
        , ("(add3 5)", "8")
        -- , ("(label prod (lambda (n) (cond ((eq n 0) 1) ('t (* n (prod (- n 1)))))))", "lambda n")
        -- , ("(prod 4)", "24")
        , ("(label Z (lambda (f) ((lambda (x) (f (lambda (v) ((x x) v)))) (lambda (x) (f (lambda (v) ((x x) v)))))))", "lambda f")
        , ("(label fib. (lambda (g) (lambda (n) (cond ((eq n 0) 0) ((eq n 1) 1) ('t (+ (g (- n 1)) (g (- n 2))))))))", "lambda g")
        , ("(label fib (Z fib.))", "lambda n")
        , ("(fib 5)", "5")
        ]
      -- 运行所有测试，获取结果和最终环境
      (results, _) = runState (mapM runTest tests) initialEnv
      -- 运行单个测试
      runTest (expr, expected) = do
        result <- sExprToValue expr  -- 求值表达式
        let resultStr = stringOfValue result  -- 转换为字符串
        return $ "测试: " ++ expr ++ 
                if resultStr == expected
                then "  [通过]"  -- 结果匹配
                else "  [失败] 期望: " ++ expected ++ " 结果: " ++ resultStr  -- 结果不匹配
  in unlines ("Running tests..." : results ++ ["Tests completed."])  -- 格式化测试输出

-- ===== REPL 功能 =====

-- 读取-求值-打印循环
repl :: IO ()
repl = do
  putStrLn "Lisp REPL (输入 exit 退出)"
  loop initialEnv  -- 开始循环，初始环境
  where
    loop env = do
      putStr "> "  -- 提示符
      hFlush stdout  -- 确保提示符显示
      input <- getLine  -- 读取用户输入
      if input == "exit"
        then putStrLn "再见!"  -- 退出命令
        else case tokenize input of  -- 处理输入
              [] -> loop env  -- 空输入继续循环
              _ -> do
                -- 求值输入并获取新环境
                let (result, newEnv) = runState (sExprToValue input) env
                putStrLn $ stringOfValue result  -- 打印结果
                loop newEnv  -- 使用新环境继续循环

-- 主函数
main :: IO ()
main = do
  putStrLn testBuiltins  -- 运行测试
  repl  -- 启动REPL
